import 'dart:async';
import 'dart:convert';

import 'package:anx_reader/utils/log/common.dart';
import 'package:langchain/langchain.dart';

class CancelableLangchainRunner {
  static const String thinkTag = '<think/>';
  StreamSubscription<ChatResult>? _subscription;

  void cancel() {
    _subscription?.cancel();
    _subscription = null;
  }

  Stream<String> stream({
    required BaseChatModel model,
    required PromptValue prompt,
  }) {
    String thinkBuffer = '';
    String answerBuffer = '';
    bool reasoningDetected = false;
    bool answerPhaseStarted = false;

    late StreamController<String> controller;
    controller = StreamController<String>(
      onListen: () {
        final source = model.stream(prompt);
        _subscription = source.listen(
          (event) {
            final rawChunk = event.output.content;
            if (rawChunk.isEmpty) {
              return;
            }

            if (_isThinkChunk(rawChunk)) {
              reasoningDetected = true;
              final cleaned = _cleanThinkChunk(rawChunk);
              if (cleaned.isNotEmpty) {
                thinkBuffer += cleaned;
              }
            } else {
              if (reasoningDetected && !answerPhaseStarted) {
                if (rawChunk.trim().isEmpty) {
                  thinkBuffer += rawChunk;
                } else {
                  answerPhaseStarted = true;
                  answerBuffer += rawChunk;
                }
              } else {
                answerBuffer += rawChunk;
              }
            }

            final aggregated = reasoningDetected
                ? '<think>${thinkBuffer.trim()}</think>\n$answerBuffer'
                : answerBuffer;

            if (!controller.isClosed) {
              controller.add(aggregated);
            }
          },
          onError: (Object error, StackTrace stackTrace) {
            if (!controller.isClosed) {
              controller.addError(error, stackTrace);
            }
          },
          onDone: () async {
            await _closeModel(model);
            if (!controller.isClosed) {
              await controller.close();
            }
            _subscription = null;
          },
          cancelOnError: false,
        );
      },
      onCancel: () async {
        await _subscription?.cancel();
        _subscription = null;
        await _closeModel(model);
        if (!controller.isClosed) {
          await controller.close();
        }
      },
    );

    return controller.stream;
  }

  Stream<String> streamAgent({
    required BaseChatModel model,
    required List<Tool> tools,
    required List<ChatMessage> history,
    required String input,
    ChatMessage? systemMessage,
    int maxIterations = 120,
  }) {
    final controller = StreamController<String>();

    Future<void>(() async {
      final parser = const ToolsAgentOutputParser();
      final toolMap = <String, Tool>{
        for (final tool in tools) tool.name: tool,
        ExceptionTool.toolName: ExceptionTool(),
      };
      final toolSpecs = tools.cast<ToolSpec>().toList(growable: false);
      final steps = <AgentStep>[];
      final timeline = <_ReasoningItem>[];
      // String? pendingThought;
      var iterations = 0;

      void emit() {
        if (controller.isClosed) return;
        controller.add(
          _composeAgentPayload(
            timeline: timeline,
          ),
        );
      }

      void appendReplyChunk(String text) {
        if (timeline.isNotEmpty &&
            timeline.last.type == _ReasoningItemType.reply) {
          timeline.last.appendReply(text);
        } else {
          timeline.add(_ReasoningItem.reply(text));
        }
      }

      List<ChatMessage> buildScratchpad() {
        final scratchpad = <ChatMessage>[];
        final seenLogs = <int>{};

        for (final step in steps) {
          for (final logMessage in step.action.messageLog) {
            final key = identityHashCode(logMessage);
            if (seenLogs.add(key)) {
              scratchpad.add(logMessage);
            }
          }

          scratchpad.add(
            ChatMessage.tool(
              toolCallId: step.action.id,
              content: step.observation,
            ),
          );
        }

        return scratchpad;
      }

      List<ChatMessage> buildConversation() {
        return <ChatMessage>[
          if (systemMessage != null) systemMessage,
          ...history,
          ChatMessage.humanText(input),
          ...buildScratchpad(),
        ];
      }

      var streamFailed = false;

      try {
        while (iterations < maxIterations && !controller.isClosed) {
          final promptMessages = buildConversation();
          if (promptMessages.isEmpty) {
            throw StateError('Agent prompt messages cannot be empty');
          }

          final prompt = PromptValue.chat(promptMessages);
          final options = model.defaultOptions.copyWith(tools: toolSpecs);

          ChatResult? aggregated;
          final completer = Completer<void>();
          _subscription = model.stream(prompt, options: options).listen(
            (chunk) {
              final normalizedChunk = _normalizeThinkChunk(chunk);

              aggregated = aggregated == null
                  ? normalizedChunk
                  : aggregated!.concat(normalizedChunk);
              final output = aggregated!.output;

              if (output.toolCalls.isEmpty) {
                final textChunk = normalizedChunk.outputAsString;
                appendReplyChunk(textChunk);
                emit();
              }
            },
            onError: (Object error, StackTrace stack) {
              streamFailed = true;
              if (!controller.isClosed) {
                controller.addError(error, stack);
              }
              if (!completer.isCompleted) {
                completer.completeError(error, stack);
              }
            },
            onDone: () {
              _subscription = null;
              if (!completer.isCompleted) {
                completer.complete();
              }
            },
            cancelOnError: true,
          );

          await completer.future;

          if (aggregated == null) {
            throw StateError('Model returned no output');
          }

          final message = aggregated!.output;
          final hydratedMessage = _hydrateToolArguments(message);
          final actions = await parser.parseChatMessage(hydratedMessage);

          // if (message.toolCalls.isNotEmpty || pendingThought != null) {
          //   // pendingThought = null;
          // }

          var shouldStop = false;
          for (final action in actions) {
            if (action is AgentFinish) {
              shouldStop = true;
              break;
            }

            final agentAction = action as AgentAction;

            final tool = toolMap[agentAction.tool];
            if (tool == null) {
              throw Exception('Tool ${agentAction.tool} not found');
            }

            final toolStep = _ToolStep(
              action: agentAction,
              status: ToolStepStatus.pending,
            );
            timeline.add(_ReasoningItem.tool(toolStep));
            emit();

            try {
              final inputJson = agentAction.toolInput;
              String? message;
              late final dynamic toolInput;
              try {
                toolInput = tool.getInputFromJson(inputJson);
              } catch (e) {
                message = 'Invalid tool input: $e';
              }
              final observation = message == null
                  ? await tool.invoke(toolInput)
                  : 'Error: $message';
              final observationText = observation.toString();
              toolStep.status = ToolStepStatus.success;
              toolStep.output = observationText;
              toolStep.observation = observationText;
              emit();
              steps.add(
                AgentStep(
                  action: agentAction,
                  observation: observationText,
                ),
              );
            } catch (error) {
              AnxLog.severe(
                  'Tool ${agentAction.tool} execution failed: $error');
              final message = error.toString();
              toolStep.status = ToolStepStatus.failed;
              toolStep.error = message;
              toolStep.observation = message;
              appendReplyChunk('Tool ${agentAction.tool} failed: $message');
              emit();
              shouldStop = true;
              break;
            }

            if (tool.returnDirect) {
              final direct = toolStep.output ?? '';
              appendReplyChunk(direct);
              emit();
              shouldStop = true;
              break;
            }
          }

          if (shouldStop) {
            break;
          }

          iterations += 1;
        }
      } catch (error, stack) {
        if (!controller.isClosed && !streamFailed) {
          controller.addError(error, stack);
        }
      } finally {
        await _subscription?.cancel();
        _subscription = null;
        await _closeModel(model);
        if (!controller.isClosed) {
          await controller.close();
        }
      }
    });

    return controller.stream;
  }

  ChatResult _normalizeThinkChunk(ChatResult chunk) {
    final content = _normalizeThinkText(chunk.output.content);
    final output =
        AIChatMessage(content: content, toolCalls: chunk.output.toolCalls);

    return ChatResult(
      output: output,
      usage: chunk.usage,
      id: chunk.id,
      finishReason: chunk.finishReason,
      metadata: chunk.metadata,
    );
  }

  String _normalizeThinkText(String text) {
    if (text.isEmpty || !_isThinkChunk(text)) {
      return text;
    }
    return _cleanThinkChunk(text);
  }

  String _composeAgentPayload({
    required List<_ReasoningItem> timeline,
  }) {
    final buffer = StringBuffer();
    for (final item in timeline) {
      final tag = item.toTag();
      if (tag.isNotEmpty) {
        buffer.write(tag);
      }
    }
    return buffer.toString();
  }

  bool _isThinkChunk(String chunk) {
    return chunk.startsWith(thinkTag);
  }

  String _cleanThinkChunk(String chunk) {
    return chunk.substring(thinkTag.length);
  }

  AIChatMessage _hydrateToolArguments(AIChatMessage message) {
    if (message.toolCalls.isEmpty) {
      return message;
    }

    var mutated = false;
    final enrichedToolCalls = <AIChatMessageToolCall>[];

    for (final toolCall in message.toolCalls) {
      if (toolCall.arguments.isNotEmpty ||
          toolCall.argumentsRaw.trim().isEmpty) {
        enrichedToolCalls.add(toolCall);
        continue;
      }

      try {
        final decoded = jsonDecode(toolCall.argumentsRaw);
        if (decoded is Map<String, dynamic>) {
          enrichedToolCalls.add(
            AIChatMessageToolCall(
              id: toolCall.id,
              name: toolCall.name,
              argumentsRaw: toolCall.argumentsRaw,
              arguments: decoded,
            ),
          );
          mutated = true;
          continue;
        }
      } catch (_) {
        // Keep original tool call if decoding fails.
      }

      enrichedToolCalls.add(toolCall);
    }

    if (!mutated) {
      return message;
    }

    return AIChatMessage(
      content: message.content,
      toolCalls: enrichedToolCalls,
    );
  }

  Future<void> _closeModel(BaseChatModel model) async {
    try {
      model.close();
    } catch (_) {
      // ignore close errors
    }
  }
}

class _ToolStep {
  _ToolStep({
    required this.action,
    required this.status,
  }) : observation = '';

  final AgentAction action;
  ToolStepStatus status;
  String observation;
  String? output;
  String? error;

  AgentStep toAgentStep() =>
      AgentStep(action: action, observation: observation);

  String toTag() {
    String? encode(String? value) {
      if (value == null || value.isEmpty) {
        return null;
      }
      final encoded = base64Encode(utf8.encode(value));
      return _escapeAttr(encoded);
    }

    final buffer = StringBuffer(
      '<tool-step name=\'${_escapeAttr(action.tool)}\' '
      "status='${status.name}'",
    );
    final inputEncoded = encode(jsonEncode(action.toolInput));
    if (inputEncoded != null) {
      buffer.write(" input_b64='$inputEncoded'");
    }
    final outputEncoded = encode(output);
    if (outputEncoded != null) {
      buffer.write(" output_b64='$outputEncoded'");
    }
    final errorEncoded = encode(error);
    if (errorEncoded != null) {
      buffer.write(" error_b64='$errorEncoded'");
    }
    buffer.write('/>');
    return buffer.toString();
  }
}

enum ToolStepStatus { pending, success, failed }

String _escapeAttr(String value) {
  return Uri.encodeComponent(value);
}

enum _ReasoningItemType { reply, tool }

class _ReasoningItem {
  _ReasoningItem.reply(String text)
      : reply = text,
        toolStep = null,
        type = _ReasoningItemType.reply;

  _ReasoningItem.tool(this.toolStep)
      : reply = null,
        type = _ReasoningItemType.tool;

  String? reply;
  final _ToolStep? toolStep;
  final _ReasoningItemType type;

  void appendReply(String text) {
    if (type != _ReasoningItemType.reply) {
      return;
    }
    reply = (reply ?? '') + text;
  }

  String toTag() {
    switch (type) {
      case _ReasoningItemType.reply:
        final text = reply;
        if (text == null || text.isEmpty) {
          return '';
        }
        final encoded = base64Encode(utf8.encode(text));
        return "<reply text_b64='${_escapeAttr(encoded)}'/>";
      case _ReasoningItemType.tool:
        if (toolStep == null) {
          return '';
        }
        return toolStep!.toTag();
    }
  }
}
